"""
传统机器学习模型训练模块
包含：随机森林、SVM等模型的训练和评估
"""

from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.model_selection import GridSearchCV, StratifiedKFold
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
import joblib
import numpy as np

class TraditionalModelTrainer:
    def __init__(self, model_type='random_forest'):
        """初始化训练器
        参数：
            model_type: 模型类型 ['random_forest', 'svm']
        """
        self.model_type = model_type
        self.models = {
            'random_forest': RandomForestClassifier(class_weight='balanced'),
            'svm': SVC(probability=True, class_weight='balanced')
        }
        # 超参数搜索空间
        self.param_grid = {
            'random_forest': {
                'n_estimators': [100, 200],  # 树的数量
                'max_depth': [10, 20]        # 最大深度
            },
            'svm': {
                'C': [0.1, 1, 10],          # 正则化参数
                'gamma': ['scale', 'auto']  # 核函数参数
            }
        }
    
    def cross_validate(self, X, y, n_splits=5):
        """分层K折交叉验证
        参数：
            X: 特征矩阵
            y: 标签
            n_splits: 折数
        返回：
            各指标的平均值字典
        """
        kf = StratifiedKFold(n_splits=n_splits)
        scores = {'accuracy': [], 'f1': [], 'auc': []}
        
        for train_idx, test_idx in kf.split(X, y):
            # 数据划分
            X_train, X_test = X[train_idx], X[test_idx]
            y_train, y_test = y[train_idx], y[test_idx]
            
            # 训练与预测
            model = self._train(X_train, y_train)
            pred = model.predict(X_test)
            proba = model.predict_proba(X_test)[:, 1]  # 预测概率
            
            # 指标计算
            scores['accuracy'].append(accuracy_score(y_test, pred))
            scores['f1'].append(f1_score(y_test, pred))
            scores['auc'].append(roc_auc_score(y_test, proba))
            
        return {k: np.mean(v) for k, v in scores.items()}
    
    def _train(self, X, y):
        """执行网格搜索训练
        返回最佳模型实例
        """
        grid_search = GridSearchCV(
            estimator=self.models[self.model_type],
            param_grid=self.param_grid[self.model_type],
            cv=3,
            n_jobs=-1,
            scoring='f1'
        )
        grid_search.fit(X, y)
        return grid_search.best_estimator_
    
    def save_model(self, model, path):
        """保存模型到文件
        使用joblib保存以支持大文件
        """
        joblib.dump(model, path, compress=3)

def train_random_forest(X_train, y_train):
    """训练随机森林分类器"""
    model = RandomForestClassifier(
        n_estimators=100,
        max_depth=10,
        random_state=42,
        class_weight='balanced'
    )
    model.fit(X_train, y_train)
    return model

def evaluate_model(model, X_test, y_test):
    """模型评估"""
    y_pred = model.predict(X_test)
    print(classification_report(y_test, y_pred))
    return y_pred

def save_model(model, file_path):
    """保存训练好的模型"""
    joblib.dump(model, file_path) 